import os
import cv2
from PIL import Image
import numpy as np


def convert_to_png(input_folder):
    """
    将文件夹内的所有图片转换为 PNG 格式，并覆盖原文件。
    :param input_folder: 输入文件夹
    """
    for filename in os.listdir(input_folder):
        input_path = os.path.join(input_folder, filename)
        if filename.lower().endswith((".jpg", ".jpeg", ".webp")):
            try:
                img = Image.open(input_path).convert("RGBA")
                output_path = os.path.splitext(input_path)[0] + ".png"
                img.save(output_path, "PNG")
                os.remove(input_path)  # 删除原文件
                print(f"已转换为 PNG: {output_path}")
            except Exception as e:
                print(f"转换失败 {input_path}: {e}")


def center_crop(input_folder, output_folder, target_size=(512, 512)):
    """
    以图片中心为参考调整图片到指定尺寸，不足部分用白色填充。
    :param input_folder: 输入文件夹路径
    :param output_folder: 输出文件夹路径
    :param target_size: 目标尺寸 (width, height)
    :return: 处理后的图片
    """
    # 确保输出文件夹存在
    if not os.path.exists(output_folder):
        os.makedirs(output_folder)

    # 遍历输入文件夹中的所有图片
    for filename in os.listdir(input_folder):
        if filename.lower().endswith((".png", ".jpg", ".jpeg")):
            image_path = os.path.join(input_folder, filename)
            img = cv2.imread(image_path, cv2.IMREAD_UNCHANGED)
            
            if img is None:
                print(f"无法读取图片: {image_path}")
                continue

            h, w = img.shape[:2]
            target_w, target_h = target_size
            
            # 计算最小缩放比例保持原图内容完整
            scale = min(target_w / w, target_h / h)
            new_w = int(w * scale)
            new_h = int(h * scale)
            
            # 缩放图片
            img = cv2.resize(img, (new_w, new_h), interpolation=cv2.INTER_AREA)
            
            # 创建新画布并用白色填充
            canvas = np.ones((target_h, target_w, 3), dtype=np.uint8) * 255
            
            # 计算放置位置
            x_offset = (target_w - new_w) // 2
            y_offset = (target_h - new_h) // 2
            
            # 将图片居中放置在新画布上
            canvas[y_offset:y_offset+new_h, x_offset:x_offset+new_w] = img
            
            # 保存处理后的图片
            output_path = os.path.join(output_folder, filename)
            cv2.imwrite(output_path, canvas)
            print(f"已保存调整后的图片: {output_path}")


def save_image(output_path, image):
    """
    仅保存为 PNG 格式，优化压缩质量。
    :param output_path: 输出文件路径
    :param image: 需要保存的图片
    """
    cv2.imwrite(output_path, image, [cv2.IMWRITE_PNG_COMPRESSION, 3])


def center_crop_pro(input_folder, output_folder, target_size=(512, 512)):
    """
    智能缩放裁剪图片，保持内容比例不变形。
    :param input_folder: 输入文件夹路径
    :param output_folder: 输出文件夹路径
    :param target_size: 目标尺寸 (width, height)
    """
    # 确保输出文件夹存在
    if not os.path.exists(output_folder):
        os.makedirs(output_folder)

    # 遍历输入文件夹中的所有图片
    for filename in os.listdir(input_folder):
        if filename.lower().endswith((".png", ".jpg", ".jpeg")):
            input_path = os.path.join(input_folder, filename)
            output_path = os.path.join(output_folder, filename)
            
            img = cv2.imread(input_path, cv2.IMREAD_UNCHANGED)
            if img is None:
                print(f"无法读取图片: {input_path}")
                continue

            h, w = img.shape[:2]
            target_w, target_h = target_size
            
            # 计算缩放比例使图片完全填满目标尺寸
            scale = max(target_w / w, target_h / h)
            new_w = int(w * scale)
            new_h = int(h * scale)
            
            # 缩放图片
            img = cv2.resize(img, (new_w, new_h), interpolation=cv2.INTER_AREA)
            
            # 计算裁剪区域
            left = max((new_w - target_w) // 2, 0)
            top = max((new_h - target_h) // 2, 0)
            right = min(left + target_w, new_w)
            bottom = min(top + target_h, new_h)

            cropped_img = img[top:bottom, left:right]
            
            # 保存裁剪后的图片
            cv2.imwrite(output_path, cropped_img)

def clear_input_folder():
    """
    清空input_images文件夹
    """
    for filename in os.listdir("./input_images"):
        file_path = os.path.join("./input_images", filename)
        try:
            if os.path.isfile(file_path):
                os.unlink(file_path)
        except Exception as e:
            print(f"删除文件{file_path}失败: {e}")
    return "input_images文件夹已清空"

def clear_output_folder():
    """
    清空output_images文件夹
    """
    for filename in os.listdir("./output_images"):
        file_path = os.path.join("./output_images", filename)
        try:
            if os.path.isfile(file_path):
                os.unlink(file_path)
        except Exception as e:
            print(f"删除文件{file_path}失败: {e}")
    return "output_images文件夹已清空"


def batch_crop_images(input_folder, output_folder, target_size=(512, 512)):
    """
    批量裁剪文件夹中的图片。
    :param input_folder: 输入文件夹
    :param output_folder: 输出文件夹
    :param target_size: 目标尺寸 (width, height)
    """
    if not os.path.exists(output_folder):
        os.makedirs(output_folder)

    for filename in os.listdir(input_folder):
        if filename.lower().endswith((".png", ".jpg", ".jpeg")):
            image_path = os.path.join(input_folder, filename)
            cropped_img = center_crop(image_path, target_size)

            if cropped_img is not None:
                output_path = os.path.join(output_folder, filename)
                save_image(output_path, cropped_img)
                print(f"已保存裁剪图片: {output_path}")


if __name__ == "__main__":
    input_folder = "./input_images"  # 输入文件夹路径
    output_folder = "./output_images"  # 输出文件夹路径
    target_size = (768, 1024)  # 目标尺寸

    # center_crop(input_folder, output_folder, target_size)  # 裁剪图片
    center_crop_pro(input_folder, output_folder, target_size)
    # batch_crop_images(input_folder, output_folder, target_size)  # 裁剪图片
